"""
机器学习策略(随机森林分类)
策略逻辑：
1. 使用技术指标作为特征
2. 预测未来N日涨跌概率
3. 当上涨概率高时买入，下跌概率高时卖出
"""

import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from backtrader import Strategy

class MLRandomForestStrategy(Strategy):
    params = (
        ('lookback', 30),  # 特征计算周期
        ('hold_period', 5),  # 持仓周期
        ('n_estimators', 100),  # 随机森林树数量
    )

    def __init__(self):
        # 初始化模型
        self.model = RandomForestClassifier(
            n_estimators=self.p.n_estimators,
            random_state=42
        )
        
        # 准备特征数据
        self.pct_change = self.data.close.pct_change()
        self.ma5 = bt.indicators.SMA(self.data.close, period=5)
        self.ma20 = bt.indicators.SMA(self.data.close, period=20)
        self.rsi = bt.indicators.RSI(self.data.close)
        self.macd = bt.indicators.MACD(self.data.close)
        
        # 训练数据缓存
        self.X = []
        self.y = []

    def next(self):
        if len(self) < self.p.lookback + self.p.hold_period:
            return
            
        # 准备特征
        features = [
            self.pct_change[0],
            (self.ma5[0] - self.ma20[0]) / self.ma20[0],
            self.rsi[0],
            self.macd.macd[0] - self.macd.signal[0]
        ]
        
        # 添加训练数据
        if len(self) % 20 == 0:  # 定期重新训练
            self.X.append(features)
            future_return = self.data.close[self.p.hold_period] / self.data.close[0] - 1
            self.y.append(1 if future_return > 0 else 0)
            
            if len(self.X) > 100:  # 足够样本后开始训练
                self.model.fit(self.X, self.y)
        
        # 预测
        if hasattr(self.model, 'predict_proba'):
            proba = self.model.predict_proba([features])[0][1]
            if proba > 0.7 and not self.position:
                self.buy()
            elif proba < 0.3 and self.position:
                self.close()

if __name__ == '__main__':
    cerebro = bt.Cerebro()
    cerebro.addstrategy(MLRandomForestStrategy)
    # 这里添加数据和其他配置
    cerebro.run()